//
// Created by hamlet on 23-3-11.
//

#include "Model"
#include <filesystem>
#include <fmt/core.h>
#include "Solver"
#include "Util"

using namespace FansionML;

namespace fs = std::filesystem;

void Model::explainData() const {
    fmt::print("数据维度：{} * {}\n标签维度：{} * 1\n标签种类：{}\n",
               static_cast<int>(data.rows()), static_cast<int>(data.cols()), static_cast<int>(label.size()),
               static_cast<int>(label_name.size()));
}

void Model::savePreprocessData(const char *path) const {
    std::ofstream out{path};
    assert(out.is_open());
    writeData(out, data);
    writeLabel(out, label);
    writeName(out, file_name);
    writeName(out, word_name);
    writeName(out, label_name);
}

void Model::loadPreprocessData(const char *path) {
    std::ifstream in{path};
    assert(in.is_open());
    readData(in, data);
    readLabel(in, label);
    readName(in, file_name);
    readName(in, word_name);
    readName(in, label_name);
}

void Model::setTestDataDir(const char *path) {
    test_data_dir = path;
}

void Model::divideTrainTestData(int n) {
    // 目前实现为了简单，没有采用完全随机划分，而是采取取余的方式
    is_part = false;
    n_fold = n;
}

void Model::divideTrainTestDataPart(int n, int i, int beg, int end) {
    is_part = true;
    n_fold = n;
    cur_fold = i;
    assign_begin = beg;
    assign_end = end;
}

std::pair<int, int> Model::getSeq() {
    auto sq = full_size / 100;
    auto i_b = sq * assign_begin;
    auto i_e = assign_end == 100 ? full_size : sq * assign_end;
    return {i_b, i_e - i_b};
}

void Model::trainAndEvaluate(Solver *solver) {
    if (solver == nullptr) return;
    std::vector<int> train_index, test_index;
    auto sz = data.rows() / n_fold;
    double right = 0;
    int i_b, i_e;
    if (is_part) {
        i_b = i_e = cur_fold;
    } else {
        i_b = 0, i_e = n_fold - 1;
    }
    for (cur_fold = i_b; cur_fold <= i_e; ++cur_fold) {
        // divide label
        train_index.reserve(data.rows() - sz);
        test_index.reserve(sz + 1);
        int tt = 0;
        for (int j = 0; j < data.rows(); ++j) {
            if (tt == cur_fold) {
                test_index.emplace_back(j);
            } else {
                train_index.emplace_back(j);
            }
            ++tt;
            if (tt == n_fold) {
                tt = 0;
            }
        }
        train_label = label(train_index);
        test_label = label(test_index);

        // divide data
        train_data = DataT{static_cast<Eigen::Index>(train_index.size()), data.cols()};
        test_data = DataT{static_cast<Eigen::Index>(test_index.size()), data.cols()};
        for (int k = 0; k < data.outerSize(); ++k) {
            for (DataT::InnerIterator it(data, k); it; ++it) {
                auto a = it.row() / n_fold, b = it.row() % n_fold;
                if (b == cur_fold) {
                    test_data.insert(a, it.col()) = it.value();
                } else {
                    auto r = it.row() - a;
                    if (b > cur_fold) --r;
                    train_data.insert(r, it.col()) = it.value();
                }
            }
        }
        train_data.makeCompressed();
        test_data.makeCompressed();

        train_index.clear();
        test_index.clear();

        full_size = test_data.rows();
        if (is_part) {
            auto sq = getSeq();
            test_data = test_data.block(sq.first, 0, sq.second, test_data.cols());
            bool ch = test_label(0) == 0;
            test_label = test_label(Eigen::seqN(sq.first, sq.second));
            if (ch) {
                for (int i = 0; i < 8; ++i) {
                    test_label(i) = 0;
                }
            }
        }

        // train and test
        solver->train(train_data, train_label);
        predict_label = solver->predictAll(test_data);

        if (!is_part) {
            // 比较正确率
            right += calcConfusionMatrix();
        }

        // save
        saveTestDataPart();
    }
    if (!is_part) {
        right /= n_fold;
        fmt::print("{}倍交叉训练平均正确率：{:.2f}%\n", n_fold, right);
    }
}

double Model::calcConfusionMatrix() const {
    auto sz = static_cast<Eigen::Index>(label_name.size());
    Eigen::MatrixXd result = Eigen::MatrixXd::Zero(sz, sz + 1);
    for (int i = 0; i < test_label.size(); ++i) {
        if (predict_label(i) >= 0) {
            ++result(test_label(i), predict_label(i));
        } else {
            ++result(test_label(i), sz);
        }
    }
    fmt::print("第{}次训练：\n", cur_fold + 1);
    int lg = 0;
    for (auto &nm: label_name) {
        if (nm.size() > lg) {
            lg = nm.size();
        }
    }
    ++lg;
    fmt::print("  ");
    fmt::print("{:>{}}", "classname", lg);
    for (int i = 0; i < sz; ++i) {
        fmt::print("{:>4}", i);
    }
    fmt::print(" unknown");
    fmt::print(" :total\n");
    Eigen::VectorXd l_nums{sz}, l_right{sz};
    for (int i = 0; i < sz; ++i) {
        fmt::print("{:>2}", i);
        fmt::print("{:>{}}", label_name[i], lg);
        for (int j = 0; j < sz; ++j) {
            if (result(i, j) == 0) {
                fmt::print("   .");
            } else {
                fmt::print("{:>4}", result(i, j));
            }
        }
        fmt::print("{:>8}", result(i, sz));
        l_nums(i) = result.row(i).sum();
        fmt::print(" :{:<5}", l_nums(i));
        l_right(i) = static_cast<double>(result(i, i)) * 100 / l_nums(i);
        fmt::print("{:<7.2f}%\n", l_right(i));
    }
    double right = l_nums.dot(l_right) / test_label.size();
    fmt::print("本次训练平均正确率：{:.2f}%\n\n", right);
    return right;
}

void Model::mergeTestResult() {
    fs::directory_entry entry{test_data_dir};
    if (!entry.exists() || !entry.is_directory()) {
        fmt::print("无法打开目录：{}\n", test_data_dir);
        exit(-1);
    }
    std::vector<bool> list_inited;
    std::vector<LabelT> list_test_label;
    std::vector<LabelT> list_predict_label;
    for (auto &p: fs::directory_iterator(test_data_dir)) {
        if (p.is_regular_file()) {
            loadTestDataPart(p.path().c_str());
            if (list_inited.empty()) {
                list_inited.resize(n_fold, false);
                list_test_label.resize(n_fold);
                list_predict_label.resize(n_fold);
            }
            if (!list_inited[cur_fold]) {
                list_inited[cur_fold] = true;
                list_test_label[cur_fold] = LabelT(full_size);
                list_predict_label[cur_fold] = LabelT(full_size);
                list_test_label[cur_fold].setConstant(-1);
                list_predict_label[cur_fold].setConstant(-1);
            }
            auto sq = getSeq();
            for (int i = 0; i < test_label.size(); ++i) {
                list_test_label[cur_fold](sq.first + i) = test_label(i);
                list_predict_label[cur_fold](sq.first + i) = predict_label(i);
            }
        }
    }
    double right = 0;
    for (cur_fold = 0; cur_fold < n_fold; ++cur_fold) {
        test_label = list_test_label[cur_fold];
        predict_label = list_predict_label[cur_fold];
        right += calcConfusionMatrix();
    }
    right /= n_fold;
    fmt::print("{}倍交叉训练平均正确率：{:.2f}%\n", n_fold, right);
}

void Model::saveTestDataPart() const {
    std::ofstream out{
            fs::path(test_data_dir) / fmt::format("test_data_{}_{}_{}_{}", n_fold, cur_fold, assign_begin, assign_end)};
    assert(out.is_open());
    out.write(reinterpret_cast<const char *>(&n_fold), sizeof(n_fold));
    out.write(reinterpret_cast<const char *>(&cur_fold), sizeof(cur_fold));
    out.write(reinterpret_cast<const char *>(&full_size), sizeof(full_size));
    out.write(reinterpret_cast<const char *>(&assign_begin), sizeof(assign_begin));
    out.write(reinterpret_cast<const char *>(&assign_end), sizeof(assign_end));
    writeLabel(out, test_label);
    writeLabel(out, predict_label);
}

void Model::loadTestDataPart(const char *path) {
    std::ifstream in{path};
    assert(in.is_open());
    in.read(reinterpret_cast<char *>(&n_fold), sizeof(n_fold));
    in.read(reinterpret_cast<char *>(&cur_fold), sizeof(cur_fold));
    in.read(reinterpret_cast<char *>(&full_size), sizeof(full_size));
    in.read(reinterpret_cast<char *>(&assign_begin), sizeof(assign_begin));
    in.read(reinterpret_cast<char *>(&assign_end), sizeof(assign_end));
    readLabel(in, test_label);
    readLabel(in, predict_label);
}
